#include <windows.h>
#include "CReleaseMacro.h"
#include "CDXPrefixSum.h"
#include <cmath>
#define FX_nThread 128

CDXPrefixSum::CDXPrefixSum(CDXDevice* pDevice) : CDXBase(pDevice)
{
	m_pCountBuf = NULL;
	m_pCountSumBuf = NULL;
	m_pPingPongBuf = NULL;
	_InitEffectLibrary();
}
CDXPrefixSum::~CDXPrefixSum()
{
	SAFE_RELEASE(m_pCountBuf);
	SAFE_RELEASE(m_pCountSumBuf);
	SAFE_RELEASE(m_pPingPongBuf);
}
HRESULT CDXPrefixSum::_InitEffectLibrary()
{
	if (m_pDevice == NULL)
		return E_FAIL;
	HRESULT hr = S_OK;
	if (m_pEffectLibrary == NULL)
	{
		m_pEffectLibrary = new CDXEffectLibrary(m_pDevice);
		hr = m_pEffectLibrary->LoadEffectFromFile(L"fx\\CDXPrefixSum.fx", "CDXPrefixSum");
	}
	return hr;
}

HRESULT CDXPrefixSum::SetCountBuffers(CDXBuffer* pCountBuf, CDXBuffer* pCountSumBuf)
{
	if (pCountBuf == NULL || pCountSumBuf == NULL)
		return E_INVALIDARG;
	ID3D11Device* pDevice = m_pDevice->GetD3D11Device();
	ID3D11DeviceContext* pDeviceContext = m_pDevice->GetD3D11DeviceContext();
	HRESULT hr = S_OK;
	SAFE_RELEASE(m_pCountBuf);
	SAFE_RELEASE(m_pCountSumBuf);
	SAFE_RELEASE(m_pPingPongBuf);
	m_pCountBuf = pCountBuf;
	m_pCountBuf->AddRef();
	m_pCountSumBuf = pCountSumBuf;
	m_pCountSumBuf->AddRef();
	
	D3D11_BUFFER_DESC desc;
	UINT szElem = 0, nElem = 0; 
	m_pCountSumBuf->GetDesc(desc, szElem, nElem);
	BOOL bCreatePingPong = TRUE;
	if (m_pPingPongBuf != NULL)
	{
		D3D11_BUFFER_DESC pingDesc;
		UINT szPingElem = 0, nPingElem = 0;
		m_pPingPongBuf->GetDesc(pingDesc, szPingElem, nPingElem);
		if (memcmp(&pingDesc, &desc, sizeof(D3D11_BUFFER_DESC) == 0) && szElem == szPingElem && nElem == nPingElem)
		{
			bCreatePingPong = FALSE;
		}
	}
	if (bCreatePingPong)
	{
		m_pPingPongBuf = new CDXBuffer(m_pDevice);
		hr = m_pPingPongBuf->Create(&desc, szElem, nElem, NULL);
	}
	return hr;
}
HRESULT CDXPrefixSum::GetCountSumBuf(CDXBuffer** ppCountSumBuf)
{
	if (m_pCountSumBuf == NULL)
		return E_FAIL;
	*ppCountSumBuf = m_pCountSumBuf;
	return S_OK;
}
HRESULT CDXPrefixSum::_PrefixSum(DXGI_FORMAT srvFmt, DXGI_FORMAT uavFmt, LPCSTR srcVarName, LPCSTR dstVarName,
								LPCSTR nameTech, UINT idxPass)
{
	if (m_pCountBuf == NULL || m_pCountSumBuf == NULL || m_pPingPongBuf == NULL || m_pDevice == NULL)
		return E_FAIL;

	HRESULT hr = S_OK;
	CDXEffect* pCDXEffect = NULL;
	m_pEffectLibrary->GetEffectByName("CDXPrefixSum", &pCDXEffect);
	ID3DX11Effect* pEffect = pCDXEffect->GetEffect();

	ID3D11DeviceContext* pDeviceContext = m_pDevice->GetD3D11DeviceContext();

	D3D11_BUFFER_DESC desc;
	UINT nElem = 0, szElem = 0;
	m_pCountSumBuf->GetDesc(desc, szElem, nElem);

	CDXBuffer* pTmpCountBuffer[2] = { m_pCountSumBuf, m_pPingPongBuf};
	
	ID3D11ShaderResourceView* pTmpCountBufferSRV[2] = { NULL, NULL};
	ID3D11UnorderedAccessView* pTmpCountBufferUAV[2] = { NULL, NULL};
	for (int i =0; i< 2; i++)
	{
		pTmpCountBuffer[i]->GetShaderResourceView(srvFmt, &(pTmpCountBufferSRV[i]));
		pTmpCountBuffer[i]->GetUnorderedAccessView(uavFmt, 0, &(pTmpCountBufferUAV[i]));
	}

	ID3DX11EffectTechnique* pTechnique = pEffect->GetTechniqueByName(nameTech);
	ID3DX11EffectPass* pPass = pTechnique->GetPassByIndex(idxPass);
	
	ID3DX11EffectScalarVariable* pCPassVar = pEffect->GetVariableByName("cPass")->AsScalar();
	ID3DX11EffectScalarVariable* pnGroupXVar = pEffect->GetVariableByName("nGroupX")->AsScalar();
	ID3DX11EffectScalarVariable* pnGroupYVar = pEffect->GetVariableByName("nGroupY")->AsScalar();
	ID3DX11EffectScalarVariable* pnGroupZVar = pEffect->GetVariableByName("nGroupZ")->AsScalar();
	ID3DX11EffectShaderResourceVariable* pV2FCountBuffer_R_Var = pEffect->GetVariableByName(srcVarName)->AsShaderResource();
	ID3DX11EffectUnorderedAccessViewVariable* pV2FCountBuffer_Var =pEffect->GetVariableByName(dstVarName)->AsUnorderedAccessView();

	UINT nPass = (UINT)ceil(log((float)nElem)/log(2.0));
	
	ID3D11ShaderResourceView *const pSRVNULL[D3D11_COMMONSHADER_INPUT_RESOURCE_SLOT_COUNT] = {NULL};
	ID3D11UnorderedAccessView* const pUAVNULL[D3D11_PS_CS_UAV_REGISTER_COUNT] = {NULL};
	ID3D11Buffer* const pCSCBNULL[D3D11_COMMONSHADER_CONSTANT_BUFFER_API_SLOT_COUNT] = {NULL};
	UINT pUAVInitCount[D3D11_PS_CS_UAV_REGISTER_COUNT] = {0};
	hr = m_pCountSumBuf->CopyFrom(m_pCountBuf);

	for (UINT cPass = 0; cPass < nPass; cPass++)
	{
		pCPassVar->SetInt(cPass);
		UINT nGroupX = (UINT)ceil((nElem - pow((float)2, (float)cPass))/FX_nThread);
		UINT nGroupY = 1;
		UINT nGroupZ = 1;
		if (nGroupX > 65535)
		{
			nGroupX = (UINT)ceil((float)sqrt((float)nGroupX));
			nGroupY = nGroupX;
		}

		void* para[] = {this, &nGroupX, &nGroupY, &nGroupZ};
		hr = pCPassVar->SetInt(cPass);
		hr = pnGroupXVar->SetInt(nGroupX);
		hr = pnGroupYVar->SetInt(nGroupY);
		hr = pnGroupZVar->SetInt(nGroupZ);
		hr = pV2FCountBuffer_R_Var->SetResource(pTmpCountBufferSRV[cPass%2]);
		hr = pV2FCountBuffer_Var->SetUnorderedAccessView(pTmpCountBufferUAV[(cPass+1)%2]);
		pPass->Apply(0, pDeviceContext);
		
		pDeviceContext->Dispatch(nGroupX, nGroupY, nGroupZ);
		
		pDeviceContext->CSSetShaderResources(0, D3D11_COMMONSHADER_INPUT_RESOURCE_SLOT_COUNT, pSRVNULL);
		pDeviceContext->CSSetUnorderedAccessViews(0, D3D11_PS_CS_UAV_REGISTER_COUNT, pUAVNULL, pUAVInitCount);
	}
	
	if ((nPass)%2)
	{
		hr = pTmpCountBuffer[0]->CopyFrom(pTmpCountBuffer[1]);
		if (FAILED(hr))
			return hr;
	}
	return hr;
}
HRESULT CDXPrefixSum::PrefixSum(CDXPrefixSum_TYPE type)
{
	HRESULT hr = S_OK;
	switch (type)
	{
	case UINT1:
		hr = _PrefixSum(DXGI_FORMAT_R32_UINT, DXGI_FORMAT_R32_UINT, "g_CountBuffer_UINT1", "g_CountSumBuffer_RW_UINT1", "Tech_PrefixSum_UINT1", 0);
		break;
	case UINT2:
		hr = _PrefixSum(DXGI_FORMAT_R32G32_UINT, DXGI_FORMAT_R32G32_UINT, "g_CountBuffer_UINT2", "g_CountSumBuffer_RW_UINT2", "Tech_PrefixSum_UINT2", 0);
		break;
	case UINT3:
		hr = _PrefixSum(DXGI_FORMAT_R32G32B32_UINT, DXGI_FORMAT_R32G32B32_UINT, "g_CountBuffer_UINT3", "g_CountSumBuffer_RW_UINT3", "Tech_PrefixSum_UINT3", 0);
		break;
	case UINT4:
		hr = _PrefixSum(DXGI_FORMAT_R32G32B32A32_UINT, DXGI_FORMAT_R32G32B32A32_UINT, "g_CountBuffer_UINT4", "g_CountSumBuffer_RW_UINT4", "Tech_PrefixSum_UINT4", 0);
		break;
	case FLOAT1:
		hr = _PrefixSum(DXGI_FORMAT_R32_FLOAT, DXGI_FORMAT_R32_FLOAT, "g_CountBuffer_FLOAT1", "g_CountSumBuffer_RW_FLOAT1", "Tech_PrefixSum_FLOAT1", 0);
		break;
	case FLOAT2:
		hr = _PrefixSum(DXGI_FORMAT_R32G32_FLOAT, DXGI_FORMAT_R32G32_FLOAT, "g_CountBuffer_FLOAT2", "g_CountSumBuffer_RW_FLOAT2", "Tech_PrefixSum_FLOAT2", 0);
		break;
	case FLOAT3:
		hr = _PrefixSum(DXGI_FORMAT_R32G32B32_FLOAT, DXGI_FORMAT_R32G32B32_FLOAT, "g_CountBuffer_FLOAT3", "g_CountSumBuffer_RW_FLOAT3", "Tech_PrefixSum_FLOAT3", 0);
		break;
	case FLOAT4:
		hr = _PrefixSum(DXGI_FORMAT_R32G32B32A32_FLOAT, DXGI_FORMAT_R32G32B32A32_FLOAT, "g_CountBuffer_FLOAT4", "g_CountSumBuffer_RW_FLOAT4", "Tech_PrefixSum_FLOAT4", 0);
		break;
	}
	
	return hr;
}